{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 67,
   "id": "expected-domestic",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "ModuleList(\n",
      "  (0): BasicLayer(\n",
      "    (blocks): ModuleList(\n",
      "      (0): SwinTransformerBlock(\n",
      "        (norm1): LayerNorm((96,), eps=1e-05, elementwise_affine=True)\n",
      "        (attn): WindowAttention(\n",
      "          (qkv): Linear(in_features=96, out_features=288, bias=True)\n",
      "          (attn_drop): Dropout(p=0.0, inplace=False)\n",
      "          (proj): Linear(in_features=96, out_features=96, bias=True)\n",
      "          (proj_drop): Dropout(p=0.0, inplace=False)\n",
      "          (softmax): Softmax(dim=-1)\n",
      "        )\n",
      "        (drop_path): Identity()\n",
      "        (norm2): LayerNorm((96,), eps=1e-05, elementwise_affine=True)\n",
      "        (mlp): Mlp(\n",
      "          (fc1): Linear(in_features=96, out_features=384, bias=True)\n",
      "          (act): GELU()\n",
      "          (fc2): Linear(in_features=384, out_features=96, bias=True)\n",
      "          (drop): Dropout(p=0.0, inplace=False)\n",
      "        )\n",
      "      )\n",
      "      (1): SwinTransformerBlock(\n",
      "        (norm1): LayerNorm((96,), eps=1e-05, elementwise_affine=True)\n",
      "        (attn): WindowAttention(\n",
      "          (qkv): Linear(in_features=96, out_features=288, bias=True)\n",
      "          (attn_drop): Dropout(p=0.0, inplace=False)\n",
      "          (proj): Linear(in_features=96, out_features=96, bias=True)\n",
      "          (proj_drop): Dropout(p=0.0, inplace=False)\n",
      "          (softmax): Softmax(dim=-1)\n",
      "        )\n",
      "        (drop_path): DropPath()\n",
      "        (norm2): LayerNorm((96,), eps=1e-05, elementwise_affine=True)\n",
      "        (mlp): Mlp(\n",
      "          (fc1): Linear(in_features=96, out_features=384, bias=True)\n",
      "          (act): GELU()\n",
      "          (fc2): Linear(in_features=384, out_features=96, bias=True)\n",
      "          (drop): Dropout(p=0.0, inplace=False)\n",
      "        )\n",
      "      )\n",
      "    )\n",
      "    (downsample): PatchMerging(\n",
      "      (reduction): Linear(in_features=384, out_features=192, bias=False)\n",
      "      (norm): LayerNorm((384,), eps=1e-05, elementwise_affine=True)\n",
      "    )\n",
      "  )\n",
      "  (1): BasicLayer(\n",
      "    (blocks): ModuleList(\n",
      "      (0): SwinTransformerBlock(\n",
      "        (norm1): LayerNorm((192,), eps=1e-05, elementwise_affine=True)\n",
      "        (attn): WindowAttention(\n",
      "          (qkv): Linear(in_features=192, out_features=576, bias=True)\n",
      "          (attn_drop): Dropout(p=0.0, inplace=False)\n",
      "          (proj): Linear(in_features=192, out_features=192, bias=True)\n",
      "          (proj_drop): Dropout(p=0.0, inplace=False)\n",
      "          (softmax): Softmax(dim=-1)\n",
      "        )\n",
      "        (drop_path): DropPath()\n",
      "        (norm2): LayerNorm((192,), eps=1e-05, elementwise_affine=True)\n",
      "        (mlp): Mlp(\n",
      "          (fc1): Linear(in_features=192, out_features=768, bias=True)\n",
      "          (act): GELU()\n",
      "          (fc2): Linear(in_features=768, out_features=192, bias=True)\n",
      "          (drop): Dropout(p=0.0, inplace=False)\n",
      "        )\n",
      "      )\n",
      "      (1): SwinTransformerBlock(\n",
      "        (norm1): LayerNorm((192,), eps=1e-05, elementwise_affine=True)\n",
      "        (attn): WindowAttention(\n",
      "          (qkv): Linear(in_features=192, out_features=576, bias=True)\n",
      "          (attn_drop): Dropout(p=0.0, inplace=False)\n",
      "          (proj): Linear(in_features=192, out_features=192, bias=True)\n",
      "          (proj_drop): Dropout(p=0.0, inplace=False)\n",
      "          (softmax): Softmax(dim=-1)\n",
      "        )\n",
      "        (drop_path): DropPath()\n",
      "        (norm2): LayerNorm((192,), eps=1e-05, elementwise_affine=True)\n",
      "        (mlp): Mlp(\n",
      "          (fc1): Linear(in_features=192, out_features=768, bias=True)\n",
      "          (act): GELU()\n",
      "          (fc2): Linear(in_features=768, out_features=192, bias=True)\n",
      "          (drop): Dropout(p=0.0, inplace=False)\n",
      "        )\n",
      "      )\n",
      "    )\n",
      "    (downsample): PatchMerging(\n",
      "      (reduction): Linear(in_features=768, out_features=384, bias=False)\n",
      "      (norm): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
      "    )\n",
      "  )\n",
      "  (2): BasicLayer(\n",
      "    (blocks): ModuleList(\n",
      "      (0): SwinTransformerBlock(\n",
      "        (norm1): LayerNorm((384,), eps=1e-05, elementwise_affine=True)\n",
      "        (attn): WindowAttention(\n",
      "          (qkv): Linear(in_features=384, out_features=1152, bias=True)\n",
      "          (attn_drop): Dropout(p=0.0, inplace=False)\n",
      "          (proj): Linear(in_features=384, out_features=384, bias=True)\n",
      "          (proj_drop): Dropout(p=0.0, inplace=False)\n",
      "          (softmax): Softmax(dim=-1)\n",
      "        )\n",
      "        (drop_path): DropPath()\n",
      "        (norm2): LayerNorm((384,), eps=1e-05, elementwise_affine=True)\n",
      "        (mlp): Mlp(\n",
      "          (fc1): Linear(in_features=384, out_features=1536, bias=True)\n",
      "          (act): GELU()\n",
      "          (fc2): Linear(in_features=1536, out_features=384, bias=True)\n",
      "          (drop): Dropout(p=0.0, inplace=False)\n",
      "        )\n",
      "      )\n",
      "      (1): SwinTransformerBlock(\n",
      "        (norm1): LayerNorm((384,), eps=1e-05, elementwise_affine=True)\n",
      "        (attn): WindowAttention(\n",
      "          (qkv): Linear(in_features=384, out_features=1152, bias=True)\n",
      "          (attn_drop): Dropout(p=0.0, inplace=False)\n",
      "          (proj): Linear(in_features=384, out_features=384, bias=True)\n",
      "          (proj_drop): Dropout(p=0.0, inplace=False)\n",
      "          (softmax): Softmax(dim=-1)\n",
      "        )\n",
      "        (drop_path): DropPath()\n",
      "        (norm2): LayerNorm((384,), eps=1e-05, elementwise_affine=True)\n",
      "        (mlp): Mlp(\n",
      "          (fc1): Linear(in_features=384, out_features=1536, bias=True)\n",
      "          (act): GELU()\n",
      "          (fc2): Linear(in_features=1536, out_features=384, bias=True)\n",
      "          (drop): Dropout(p=0.0, inplace=False)\n",
      "        )\n",
      "      )\n",
      "      (2): SwinTransformerBlock(\n",
      "        (norm1): LayerNorm((384,), eps=1e-05, elementwise_affine=True)\n",
      "        (attn): WindowAttention(\n",
      "          (qkv): Linear(in_features=384, out_features=1152, bias=True)\n",
      "          (attn_drop): Dropout(p=0.0, inplace=False)\n",
      "          (proj): Linear(in_features=384, out_features=384, bias=True)\n",
      "          (proj_drop): Dropout(p=0.0, inplace=False)\n",
      "          (softmax): Softmax(dim=-1)\n",
      "        )\n",
      "        (drop_path): DropPath()\n",
      "        (norm2): LayerNorm((384,), eps=1e-05, elementwise_affine=True)\n",
      "        (mlp): Mlp(\n",
      "          (fc1): Linear(in_features=384, out_features=1536, bias=True)\n",
      "          (act): GELU()\n",
      "          (fc2): Linear(in_features=1536, out_features=384, bias=True)\n",
      "          (drop): Dropout(p=0.0, inplace=False)\n",
      "        )\n",
      "      )\n",
      "      (3): SwinTransformerBlock(\n",
      "        (norm1): LayerNorm((384,), eps=1e-05, elementwise_affine=True)\n",
      "        (attn): WindowAttention(\n",
      "          (qkv): Linear(in_features=384, out_features=1152, bias=True)\n",
      "          (attn_drop): Dropout(p=0.0, inplace=False)\n",
      "          (proj): Linear(in_features=384, out_features=384, bias=True)\n",
      "          (proj_drop): Dropout(p=0.0, inplace=False)\n",
      "          (softmax): Softmax(dim=-1)\n",
      "        )\n",
      "        (drop_path): DropPath()\n",
      "        (norm2): LayerNorm((384,), eps=1e-05, elementwise_affine=True)\n",
      "        (mlp): Mlp(\n",
      "          (fc1): Linear(in_features=384, out_features=1536, bias=True)\n",
      "          (act): GELU()\n",
      "          (fc2): Linear(in_features=1536, out_features=384, bias=True)\n",
      "          (drop): Dropout(p=0.0, inplace=False)\n",
      "        )\n",
      "      )\n",
      "      (4): SwinTransformerBlock(\n",
      "        (norm1): LayerNorm((384,), eps=1e-05, elementwise_affine=True)\n",
      "        (attn): WindowAttention(\n",
      "          (qkv): Linear(in_features=384, out_features=1152, bias=True)\n",
      "          (attn_drop): Dropout(p=0.0, inplace=False)\n",
      "          (proj): Linear(in_features=384, out_features=384, bias=True)\n",
      "          (proj_drop): Dropout(p=0.0, inplace=False)\n",
      "          (softmax): Softmax(dim=-1)\n",
      "        )\n",
      "        (drop_path): DropPath()\n",
      "        (norm2): LayerNorm((384,), eps=1e-05, elementwise_affine=True)\n",
      "        (mlp): Mlp(\n",
      "          (fc1): Linear(in_features=384, out_features=1536, bias=True)\n",
      "          (act): GELU()\n",
      "          (fc2): Linear(in_features=1536, out_features=384, bias=True)\n",
      "          (drop): Dropout(p=0.0, inplace=False)\n",
      "        )\n",
      "      )\n",
      "      (5): SwinTransformerBlock(\n",
      "        (norm1): LayerNorm((384,), eps=1e-05, elementwise_affine=True)\n",
      "        (attn): WindowAttention(\n",
      "          (qkv): Linear(in_features=384, out_features=1152, bias=True)\n",
      "          (attn_drop): Dropout(p=0.0, inplace=False)\n",
      "          (proj): Linear(in_features=384, out_features=384, bias=True)\n",
      "          (proj_drop): Dropout(p=0.0, inplace=False)\n",
      "          (softmax): Softmax(dim=-1)\n",
      "        )\n",
      "        (drop_path): DropPath()\n",
      "        (norm2): LayerNorm((384,), eps=1e-05, elementwise_affine=True)\n",
      "        (mlp): Mlp(\n",
      "          (fc1): Linear(in_features=384, out_features=1536, bias=True)\n",
      "          (act): GELU()\n",
      "          (fc2): Linear(in_features=1536, out_features=384, bias=True)\n",
      "          (drop): Dropout(p=0.0, inplace=False)\n",
      "        )\n",
      "      )\n",
      "      (6): SwinTransformerBlock(\n",
      "        (norm1): LayerNorm((384,), eps=1e-05, elementwise_affine=True)\n",
      "        (attn): WindowAttention(\n",
      "          (qkv): Linear(in_features=384, out_features=1152, bias=True)\n",
      "          (attn_drop): Dropout(p=0.0, inplace=False)\n",
      "          (proj): Linear(in_features=384, out_features=384, bias=True)\n",
      "          (proj_drop): Dropout(p=0.0, inplace=False)\n",
      "          (softmax): Softmax(dim=-1)\n",
      "        )\n",
      "        (drop_path): DropPath()\n",
      "        (norm2): LayerNorm((384,), eps=1e-05, elementwise_affine=True)\n",
      "        (mlp): Mlp(\n",
      "          (fc1): Linear(in_features=384, out_features=1536, bias=True)\n",
      "          (act): GELU()\n",
      "          (fc2): Linear(in_features=1536, out_features=384, bias=True)\n",
      "          (drop): Dropout(p=0.0, inplace=False)\n",
      "        )\n",
      "      )\n",
      "      (7): SwinTransformerBlock(\n",
      "        (norm1): LayerNorm((384,), eps=1e-05, elementwise_affine=True)\n",
      "        (attn): WindowAttention(\n",
      "          (qkv): Linear(in_features=384, out_features=1152, bias=True)\n",
      "          (attn_drop): Dropout(p=0.0, inplace=False)\n",
      "          (proj): Linear(in_features=384, out_features=384, bias=True)\n",
      "          (proj_drop): Dropout(p=0.0, inplace=False)\n",
      "          (softmax): Softmax(dim=-1)\n",
      "        )\n",
      "        (drop_path): DropPath()\n",
      "        (norm2): LayerNorm((384,), eps=1e-05, elementwise_affine=True)\n",
      "        (mlp): Mlp(\n",
      "          (fc1): Linear(in_features=384, out_features=1536, bias=True)\n",
      "          (act): GELU()\n",
      "          (fc2): Linear(in_features=1536, out_features=384, bias=True)\n",
      "          (drop): Dropout(p=0.0, inplace=False)\n",
      "        )\n",
      "      )\n",
      "      (8): SwinTransformerBlock(\n",
      "        (norm1): LayerNorm((384,), eps=1e-05, elementwise_affine=True)\n",
      "        (attn): WindowAttention(\n",
      "          (qkv): Linear(in_features=384, out_features=1152, bias=True)\n",
      "          (attn_drop): Dropout(p=0.0, inplace=False)\n",
      "          (proj): Linear(in_features=384, out_features=384, bias=True)\n",
      "          (proj_drop): Dropout(p=0.0, inplace=False)\n",
      "          (softmax): Softmax(dim=-1)\n",
      "        )\n",
      "        (drop_path): DropPath()\n",
      "        (norm2): LayerNorm((384,), eps=1e-05, elementwise_affine=True)\n",
      "        (mlp): Mlp(\n",
      "          (fc1): Linear(in_features=384, out_features=1536, bias=True)\n",
      "          (act): GELU()\n",
      "          (fc2): Linear(in_features=1536, out_features=384, bias=True)\n",
      "          (drop): Dropout(p=0.0, inplace=False)\n",
      "        )\n",
      "      )\n",
      "      (9): SwinTransformerBlock(\n",
      "        (norm1): LayerNorm((384,), eps=1e-05, elementwise_affine=True)\n",
      "        (attn): WindowAttention(\n",
      "          (qkv): Linear(in_features=384, out_features=1152, bias=True)\n",
      "          (attn_drop): Dropout(p=0.0, inplace=False)\n",
      "          (proj): Linear(in_features=384, out_features=384, bias=True)\n",
      "          (proj_drop): Dropout(p=0.0, inplace=False)\n",
      "          (softmax): Softmax(dim=-1)\n",
      "        )\n",
      "        (drop_path): DropPath()\n",
      "        (norm2): LayerNorm((384,), eps=1e-05, elementwise_affine=True)\n",
      "        (mlp): Mlp(\n",
      "          (fc1): Linear(in_features=384, out_features=1536, bias=True)\n",
      "          (act): GELU()\n",
      "          (fc2): Linear(in_features=1536, out_features=384, bias=True)\n",
      "          (drop): Dropout(p=0.0, inplace=False)\n",
      "        )\n",
      "      )\n",
      "      (10): SwinTransformerBlock(\n",
      "        (norm1): LayerNorm((384,), eps=1e-05, elementwise_affine=True)\n",
      "        (attn): WindowAttention(\n",
      "          (qkv): Linear(in_features=384, out_features=1152, bias=True)\n",
      "          (attn_drop): Dropout(p=0.0, inplace=False)\n",
      "          (proj): Linear(in_features=384, out_features=384, bias=True)\n",
      "          (proj_drop): Dropout(p=0.0, inplace=False)\n",
      "          (softmax): Softmax(dim=-1)\n",
      "        )\n",
      "        (drop_path): DropPath()\n",
      "        (norm2): LayerNorm((384,), eps=1e-05, elementwise_affine=True)\n",
      "        (mlp): Mlp(\n",
      "          (fc1): Linear(in_features=384, out_features=1536, bias=True)\n",
      "          (act): GELU()\n",
      "          (fc2): Linear(in_features=1536, out_features=384, bias=True)\n",
      "          (drop): Dropout(p=0.0, inplace=False)\n",
      "        )\n",
      "      )\n",
      "      (11): SwinTransformerBlock(\n",
      "        (norm1): LayerNorm((384,), eps=1e-05, elementwise_affine=True)\n",
      "        (attn): WindowAttention(\n",
      "          (qkv): Linear(in_features=384, out_features=1152, bias=True)\n",
      "          (attn_drop): Dropout(p=0.0, inplace=False)\n",
      "          (proj): Linear(in_features=384, out_features=384, bias=True)\n",
      "          (proj_drop): Dropout(p=0.0, inplace=False)\n",
      "          (softmax): Softmax(dim=-1)\n",
      "        )\n",
      "        (drop_path): DropPath()\n",
      "        (norm2): LayerNorm((384,), eps=1e-05, elementwise_affine=True)\n",
      "        (mlp): Mlp(\n",
      "          (fc1): Linear(in_features=384, out_features=1536, bias=True)\n",
      "          (act): GELU()\n",
      "          (fc2): Linear(in_features=1536, out_features=384, bias=True)\n",
      "          (drop): Dropout(p=0.0, inplace=False)\n",
      "        )\n",
      "      )\n",
      "      (12): SwinTransformerBlock(\n",
      "        (norm1): LayerNorm((384,), eps=1e-05, elementwise_affine=True)\n",
      "        (attn): WindowAttention(\n",
      "          (qkv): Linear(in_features=384, out_features=1152, bias=True)\n",
      "          (attn_drop): Dropout(p=0.0, inplace=False)\n",
      "          (proj): Linear(in_features=384, out_features=384, bias=True)\n",
      "          (proj_drop): Dropout(p=0.0, inplace=False)\n",
      "          (softmax): Softmax(dim=-1)\n",
      "        )\n",
      "        (drop_path): DropPath()\n",
      "        (norm2): LayerNorm((384,), eps=1e-05, elementwise_affine=True)\n",
      "        (mlp): Mlp(\n",
      "          (fc1): Linear(in_features=384, out_features=1536, bias=True)\n",
      "          (act): GELU()\n",
      "          (fc2): Linear(in_features=1536, out_features=384, bias=True)\n",
      "          (drop): Dropout(p=0.0, inplace=False)\n",
      "        )\n",
      "      )\n",
      "      (13): SwinTransformerBlock(\n",
      "        (norm1): LayerNorm((384,), eps=1e-05, elementwise_affine=True)\n",
      "        (attn): WindowAttention(\n",
      "          (qkv): Linear(in_features=384, out_features=1152, bias=True)\n",
      "          (attn_drop): Dropout(p=0.0, inplace=False)\n",
      "          (proj): Linear(in_features=384, out_features=384, bias=True)\n",
      "          (proj_drop): Dropout(p=0.0, inplace=False)\n",
      "          (softmax): Softmax(dim=-1)\n",
      "        )\n",
      "        (drop_path): DropPath()\n",
      "        (norm2): LayerNorm((384,), eps=1e-05, elementwise_affine=True)\n",
      "        (mlp): Mlp(\n",
      "          (fc1): Linear(in_features=384, out_features=1536, bias=True)\n",
      "          (act): GELU()\n",
      "          (fc2): Linear(in_features=1536, out_features=384, bias=True)\n",
      "          (drop): Dropout(p=0.0, inplace=False)\n",
      "        )\n",
      "      )\n",
      "      (14): SwinTransformerBlock(\n",
      "        (norm1): LayerNorm((384,), eps=1e-05, elementwise_affine=True)\n",
      "        (attn): WindowAttention(\n",
      "          (qkv): Linear(in_features=384, out_features=1152, bias=True)\n",
      "          (attn_drop): Dropout(p=0.0, inplace=False)\n",
      "          (proj): Linear(in_features=384, out_features=384, bias=True)\n",
      "          (proj_drop): Dropout(p=0.0, inplace=False)\n",
      "          (softmax): Softmax(dim=-1)\n",
      "        )\n",
      "        (drop_path): DropPath()\n",
      "        (norm2): LayerNorm((384,), eps=1e-05, elementwise_affine=True)\n",
      "        (mlp): Mlp(\n",
      "          (fc1): Linear(in_features=384, out_features=1536, bias=True)\n",
      "          (act): GELU()\n",
      "          (fc2): Linear(in_features=1536, out_features=384, bias=True)\n",
      "          (drop): Dropout(p=0.0, inplace=False)\n",
      "        )\n",
      "      )\n",
      "      (15): SwinTransformerBlock(\n",
      "        (norm1): LayerNorm((384,), eps=1e-05, elementwise_affine=True)\n",
      "        (attn): WindowAttention(\n",
      "          (qkv): Linear(in_features=384, out_features=1152, bias=True)\n",
      "          (attn_drop): Dropout(p=0.0, inplace=False)\n",
      "          (proj): Linear(in_features=384, out_features=384, bias=True)\n",
      "          (proj_drop): Dropout(p=0.0, inplace=False)\n",
      "          (softmax): Softmax(dim=-1)\n",
      "        )\n",
      "        (drop_path): DropPath()\n",
      "        (norm2): LayerNorm((384,), eps=1e-05, elementwise_affine=True)\n",
      "        (mlp): Mlp(\n",
      "          (fc1): Linear(in_features=384, out_features=1536, bias=True)\n",
      "          (act): GELU()\n",
      "          (fc2): Linear(in_features=1536, out_features=384, bias=True)\n",
      "          (drop): Dropout(p=0.0, inplace=False)\n",
      "        )\n",
      "      )\n",
      "      (16): SwinTransformerBlock(\n",
      "        (norm1): LayerNorm((384,), eps=1e-05, elementwise_affine=True)\n",
      "        (attn): WindowAttention(\n",
      "          (qkv): Linear(in_features=384, out_features=1152, bias=True)\n",
      "          (attn_drop): Dropout(p=0.0, inplace=False)\n",
      "          (proj): Linear(in_features=384, out_features=384, bias=True)\n",
      "          (proj_drop): Dropout(p=0.0, inplace=False)\n",
      "          (softmax): Softmax(dim=-1)\n",
      "        )\n",
      "        (drop_path): DropPath()\n",
      "        (norm2): LayerNorm((384,), eps=1e-05, elementwise_affine=True)\n",
      "        (mlp): Mlp(\n",
      "          (fc1): Linear(in_features=384, out_features=1536, bias=True)\n",
      "          (act): GELU()\n",
      "          (fc2): Linear(in_features=1536, out_features=384, bias=True)\n",
      "          (drop): Dropout(p=0.0, inplace=False)\n",
      "        )\n",
      "      )\n",
      "      (17): SwinTransformerBlock(\n",
      "        (norm1): LayerNorm((384,), eps=1e-05, elementwise_affine=True)\n",
      "        (attn): WindowAttention(\n",
      "          (qkv): Linear(in_features=384, out_features=1152, bias=True)\n",
      "          (attn_drop): Dropout(p=0.0, inplace=False)\n",
      "          (proj): Linear(in_features=384, out_features=384, bias=True)\n",
      "          (proj_drop): Dropout(p=0.0, inplace=False)\n",
      "          (softmax): Softmax(dim=-1)\n",
      "        )\n",
      "        (drop_path): DropPath()\n",
      "        (norm2): LayerNorm((384,), eps=1e-05, elementwise_affine=True)\n",
      "        (mlp): Mlp(\n",
      "          (fc1): Linear(in_features=384, out_features=1536, bias=True)\n",
      "          (act): GELU()\n",
      "          (fc2): Linear(in_features=1536, out_features=384, bias=True)\n",
      "          (drop): Dropout(p=0.0, inplace=False)\n",
      "        )\n",
      "      )\n",
      "    )\n",
      "    (downsample): PatchMerging(\n",
      "      (reduction): Linear(in_features=1536, out_features=768, bias=False)\n",
      "      (norm): LayerNorm((1536,), eps=1e-05, elementwise_affine=True)\n",
      "    )\n",
      "  )\n",
      "  (3): BasicLayer(\n",
      "    (blocks): ModuleList(\n",
      "      (0): SwinTransformerBlock(\n",
      "        (norm1): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
      "        (attn): WindowAttention(\n",
      "          (qkv): Linear(in_features=768, out_features=2304, bias=True)\n",
      "          (attn_drop): Dropout(p=0.0, inplace=False)\n",
      "          (proj): Linear(in_features=768, out_features=768, bias=True)\n",
      "          (proj_drop): Dropout(p=0.0, inplace=False)\n",
      "          (softmax): Softmax(dim=-1)\n",
      "        )\n",
      "        (drop_path): DropPath()\n",
      "        (norm2): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
      "        (mlp): Mlp(\n",
      "          (fc1): Linear(in_features=768, out_features=3072, bias=True)\n",
      "          (act): GELU()\n",
      "          (fc2): Linear(in_features=3072, out_features=768, bias=True)\n",
      "          (drop): Dropout(p=0.0, inplace=False)\n",
      "        )\n",
      "      )\n",
      "      (1): SwinTransformerBlock(\n",
      "        (norm1): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
      "        (attn): WindowAttention(\n",
      "          (qkv): Linear(in_features=768, out_features=2304, bias=True)\n",
      "          (attn_drop): Dropout(p=0.0, inplace=False)\n",
      "          (proj): Linear(in_features=768, out_features=768, bias=True)\n",
      "          (proj_drop): Dropout(p=0.0, inplace=False)\n",
      "          (softmax): Softmax(dim=-1)\n",
      "        )\n",
      "        (drop_path): DropPath()\n",
      "        (norm2): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
      "        (mlp): Mlp(\n",
      "          (fc1): Linear(in_features=768, out_features=3072, bias=True)\n",
      "          (act): GELU()\n",
      "          (fc2): Linear(in_features=3072, out_features=768, bias=True)\n",
      "          (drop): Dropout(p=0.0, inplace=False)\n",
      "        )\n",
      "      )\n",
      "    )\n",
      "  )\n",
      ")\n"
     ]
    },
    {
     "ename": "IndexError",
     "evalue": "list index out of range",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mIndexError\u001b[0m                                Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-67-56d3ea1c8c35>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m    125\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    126\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 127\u001b[0;31m \u001b[0mconv_features\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mconv_features\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;36m0\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    128\u001b[0m \u001b[0;31m# print(conv_features.shape)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    129\u001b[0m \u001b[0;31m# conv_features = torch.sigmoid(conv_features)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mIndexError\u001b[0m: list index out of range"
     ]
    }
   ],
   "source": [
    "import math\n",
    "from PIL import Image\n",
    "import matplotlib.pyplot as plt\n",
    "import torch\n",
    "from torch import nn\n",
    "import torchvision.transforms as T\n",
    "import torch.nn.functional as F\n",
    "import sys\n",
    "import matplotlib.cm as cm\n",
    "import numpy as np\n",
    "import cv2\n",
    "# sys.path.append(\"..\")\n",
    "# from nets.dlanet_wh  import DlaNet\n",
    "# from nets.resnet_inp import get_pose_net\n",
    "import os \n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = '1'\n",
    "############################第一步：import想要看特征图的网络############################\n",
    "\n",
    "# from nets.resnet_concat_vit import get_pose_net as get_pose_net1\n",
    "\n",
    "######################################################################################\n",
    "\n",
    "\n",
    "############################需要用到的函数##########################\n",
    "torch.set_grad_enabled(False)\n",
    "CLASSES = ['N/A','ship']\n",
    "COLORS = [[0.000, 0.447, 0.741], [0.850, 0.325, 0.098], [0.929, 0.694, 0.125],\n",
    "          [0.494, 0.184, 0.556], [0.466, 0.674, 0.188], [0.301, 0.745, 0.933]]\n",
    "transform = T.Compose([\n",
    "    T.ToPILImage(),\n",
    "    T.Resize([800,800]),\n",
    "    T.ToTensor(),\n",
    "    T.Normalize([0.5194416012442385,0.5378052387430711,0.533462090585746], [0.3001546018824507, 0.28620901391179554, 0.3014112676161966])\n",
    "])\n",
    "\n",
    "# def _nms(heat, kernel=10):\n",
    "#     hmax = F.max_pool2d(heat, kernel, stride=1, padding=(kernel - 1) // 2)\n",
    "#     keep = (hmax == heat).float()\n",
    "#     return heat * keep\n",
    "# def _nms_me(heat, kernel=[3,5,7,9,11]):\n",
    "#     hmax_3 = F.max_pool2d(heat, kernel[0], stride=1, padding=(kernel[0] - 1) // 2)\n",
    "#     hmax_5 = F.max_pool2d(heat, kernel[1], stride=1, padding=(kernel[1] - 1) // 2)\n",
    "#     hmax_7 = F.max_pool2d(heat, kernel[2], stride=1, padding=(kernel[2] - 1) // 2)\n",
    "# #     hmax_9 = F.max_pool2d(heat, kernel[3], stride=1, padding=(kernel[3] - 1) // 2)\n",
    "# #     hmax_11 = F.max_pool2d(heat, kernel[4], stride=1, padding=(kernel[4] - 1) // 2)\n",
    "#     hmax = torch.cat([hmax_3, hmax_5, hmax_7], 1)\n",
    "#     print(hmax.shape)\n",
    "#     keep = torch.eq(hmax, heat).float()\n",
    "#     print(keep.shape)\n",
    "#     print((heat*keep).shape)\n",
    "#     return heat * keep\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "############################第二步：打开想要看的图片,并转化成网络的输入格式############################\n",
    "im = cv2.imread(\"../../_DATASET/HRSID/JPEGImages/P0111_5400_6200_9600_10400.jpg\")\n",
    "\n",
    "\n",
    "img = transform(im).unsqueeze(0)\n",
    "#####################################################################################\n",
    "\n",
    "sys.path.append(r'../backbone')\n",
    "# from resnet import ResNet  \n",
    "# from resnet_dcn import ResNet\n",
    "# from dlanet import DlaNet\n",
    "# from dlanet import DlaNet\n",
    "from swinT import SwinTransformer\n",
    "#000748\n",
    "head = {'hm': 1, 'wh': 2, 'ang':1, 'reg': 2}\n",
    "model = SwinTransformer(heads = head,\n",
    "                 pretrain_img_size=224,\n",
    "                 patch_size=4,\n",
    "                 in_chans=3,\n",
    "                 embed_dim=96,\n",
    "                 depths=[2, 2, 18, 2],#[2, 2, 18, 2]for small,[2, 2, 6, 2]for tiny\n",
    "                 num_heads=[3, 6, 12, 24],\n",
    "                 window_size=7,\n",
    "                 mlp_ratio=4.,\n",
    "                 qkv_bias=True,\n",
    "                 qk_scale=None,\n",
    "                 drop_rate=0.,\n",
    "                 attn_drop_rate=0.,\n",
    "                 drop_path_rate=0.2,\n",
    "                 norm_layer=nn.LayerNorm,\n",
    "                 ape=False,\n",
    "                 patch_norm=True,\n",
    "                 out_indices=(0, 1, 2, 3),\n",
    "                 frozen_stages=-1,\n",
    "                 use_checkpoint=False,)\n",
    "\n",
    "# from nets.dlanet_cor_mid import DlaNet\n",
    "# model = DlaNet(num_classes=1,head_conv=256)#试一下256效果怎么样\n",
    "device = torch.device(\"cpu\")\n",
    "\n",
    "state_dict = torch.load('../../_Weights/rcenternet/HRSID/swinT_small/best.pth',map_location='cpu')  \n",
    "model.load_state_dict(state_dict)\n",
    "model.to(device)\n",
    "model.eval()\n",
    "# print(state_dict.keys())\n",
    "\n",
    "conv_features, enc_attn_weights, dec_attn_weights = [], [], [] \n",
    "\n",
    "\n",
    "#extract_feature = model.hm[-1]\n",
    "extract_feature = model.layers\n",
    "#.dec_c2\n",
    "# extract_feature = model.cor_att[-1]\n",
    "\n",
    "print(extract_feature)\n",
    "\n",
    "\n",
    "# print(type(model.hmap[-4]))\n",
    "# extract_feature = model.cor_att[-1]\n",
    "#print(model)\n",
    "# v_feature = model.v\n",
    "\n",
    "hooks = [extract_feature.register_forward_hook(lambda self, input, output: conv_features.append(output))]\n",
    "# propagate through the model \n",
    "outputs = model(img) \n",
    "\n",
    "for hook in hooks:    \n",
    "    hook.remove() \n",
    "    # don't need the list anymore\n",
    "    \n",
    "\n",
    "conv_features = conv_features[0]\n",
    "# print(conv_features.shape)\n",
    "# conv_features = torch.sigmoid(conv_features)\n",
    "# conv_features = _nms(conv_features,kernel=3)\n",
    "\n",
    "# print(conv_features)\n",
    "def draw_features(width,height,x,savename):\n",
    "#     tic=time.time()\n",
    "    fig = plt.figure(figsize=(16,16))\n",
    "    fig.subplots_adjust(left=0.05, right=0.95, bottom=0.05, top=0.95, wspace=0.05, hspace=0.05)\n",
    "    for i in range(width*height):\n",
    "        plt.subplot(height,width, i + 1)\n",
    "        plt.axis('off')\n",
    "        # plt.tight_layout()\n",
    "        img = x[0, i, :, :]\n",
    "        pmin = np.min(img)\n",
    "        pmax = np.max(img)\n",
    "        img = (img - pmin) / (pmax - pmin + 0.000001)*255\n",
    "#         print(img)\n",
    "        img=img.astype(np.uint8)  #转成unit8\n",
    "        #PIL转从cv2\n",
    "        img = cv2.cvtColor(np.asarray(img),cv2.COLOR_RGB2BGR)\n",
    "        img= cv2.applyColorMap(img, cv2.COLORMAP_JET) #生成heat map\n",
    "        img = img[:, :, ::-1]#注意cv2（BGR）和matplotlib(RGB)通道是相反的\n",
    "#        plt.figure(figsize=(10,10))\n",
    "        plt.imshow(img)\n",
    "#         plt.show()\n",
    "#         print(\"{}/{}\".format(i,width*height))\n",
    "    plt.show()\n",
    "    fig.savefig(savename, dpi=100)\n",
    "    fig.clf()\n",
    "    plt.close()\n",
    "#     print(\"time:{}\".format(time.time()-tic))\n",
    "print(len(conv_features))\n",
    "print(conv_features.shape)\n",
    "#显示原图\n",
    "fig = plt.figure(figsize=(8, 8))\n",
    "plt.axis('off')\n",
    "plt.imshow(im)\n",
    "plt.show()\n",
    "fig.clf()\n",
    "plt.close()\n",
    "\n",
    "\n",
    "#draw_features(8,12,conv_features.cpu().numpy(),\"{}/f1_conv1.png\".format(\"feature\"))   \n",
    "draw_features(4,4,conv_features.cpu().numpy(),\"{}/f1_conv1.png\".format(\"feature\")) \n",
    "\n",
    "# print(conv_features.shape)\n",
    "\n",
    "# print(conv_features.shape)\n",
    "#自定义的超参数\n",
    "# h, w = 32,32\n",
    "# im = im.resize((512,512))\n",
    "\n",
    "# fig, axs = plt.subplots(ncols=8, nrows=2, figsize=(22, 7))\n",
    "# print(axs)\n",
    "# colors = COLORS * 100\n",
    "# for idx, ax_i in enumerate(axs.T):    \n",
    "#     ax = ax_i[0]    \n",
    "#     ax.imshow(conv_features[0, idx].view(h, w),cmap='viridis')    \n",
    "# #     ax.axis('off')    \n",
    "#     ax.set_title(f'query id: {idx}')    \n",
    "#     ax = ax_i[1]    \n",
    "#     ax.imshow(im)    \n",
    "#     ax.add_patch(plt.Rectangle((xmin, ymin), xmax - xmin, ymax - ymin,fill=False, color='red', linewidth=3))\n",
    "#     ax.axis('off')    \n",
    "#     ax.set_title(CLASSES[probas[idx].argmax()])\n",
    "# fig.tight_layout()\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "pediatric-association",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "billion-render",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
